"""
PostgreSQL备份增强模块：压缩/重删/加密
包含三个独立模块：CompressionModule、DeduplicationModule、EncryptionModule
"""

import os
import zlib
import lzma
import gzip
import sqlite3
import hashlib
import json
import logging
from concurrent.futures import ThreadPoolExecutor
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes

LOGGER = logging.getLogger(__name__)

# ======================
# 压缩模块
# ======================
class CompressionModule:
    """
    数据压缩模块
    功能：
    1. 支持gzip/lzma/zlib三种压缩算法
    2. 自动计算压缩率并记录日志
    
    使用示例：
    >>> compressor = CompressionModule('gzip')
    >>> compressed = compressor.compress(b'raw_data')
    >>> decompressed = compressor.decompress(compressed)
    """
    
    def __init__(self, algorithm='gzip'):
        """
        初始化压缩模块
        :param algorithm: 压缩算法，可选gzip/lzma/zlib
        """
        self.algorithm = algorithm
        LOGGER.info(f"Initialized compression module with algorithm: {algorithm}")

    def compress(self, data):
        """
        压缩数据
        :param data: 原始字节数据
        :return: 压缩后的字节数据
        """
        original_size = len(data)
        
        # 选择压缩算法
        if self.algorithm == 'gzip':
            compressed = gzip.compress(data)
        elif self.algorithm == 'lzma':
            compressed = lzma.compress(data)
        elif self.algorithm == 'zlib':
            compressed = zlib.compress(data)
        else:
            raise ValueError(f"Unsupported compression algorithm: {self.algorithm}")

        # 计算压缩率
        compressed_size = len(compressed)
        compression_rate = (original_size - compressed_size) / original_size * 100
        LOGGER.info(f"Compression completed | Algorithm: {self.algorithm} | "
                    f"Original: {original_size} bytes | "
                    f"Compressed: {compressed_size} bytes | "
                    f"Rate: {compression_rate:.2f}%")
        
        return compressed

    def decompress(self, data):
        """
        解压数据
        :param data: 压缩后的字节数据
        :return: 原始字节数据
        """
        try:
            if self.algorithm == 'gzip':
                return gzip.decompress(data)
            elif self.algorithm == 'lzma':
                return lzma.decompress(data)
            elif self.algorithm == 'zlib':
                return zlib.decompress(data)
        except Exception as e:
            LOGGER.error(f"Decompression failed: {str(e)}")
            raise

# ======================
# 重删模块
# ======================
class DeduplicationModule:
    """
    重复数据删除模块
    功能：
    1. 基于内容分块的全局重删
    2. 使用SHA-256哈希标识数据块
    3. 自动计算重删率并记录日志
    
    使用示例：
    >>> dedup = DeduplicationModule('/var/db/block_index.db')
    >>> blocks = dedup.split_into_blocks(data)
    >>> unique_blocks, dedup_rate = dedup.process_blocks(blocks)
    """
    
    def __init__(self, index_db_path, block_size=4*1024*1024):
        """
        初始化重删模块
        :param index_db_path: 块索引数据库路径
        :param block_size: 分块大小（默认4MB）
        """
        self.block_size = block_size
        self.conn = sqlite3.connect(index_db_path)
        self._init_db()
        LOGGER.info(f"Initialized deduplication module | Block size: {block_size}")

    def _init_db(self):
        """初始化数据库表结构"""
        self.conn.execute('''
            CREATE TABLE IF NOT EXISTS blocks (
                hash TEXT PRIMARY KEY,  -- 块哈希
                size INTEGER,           -- 块大小
                ref_count INTEGER DEFAULT 1 -- 引用计数
            )
        ''')
        self.conn.commit()

    def split_into_blocks(self, data):
        """
        将数据分割为固定大小的块
        :param data: 输入数据
        :return: 块列表
        """
        return [data[i:i+self.block_size] 
               for i in range(0, len(data), self.block_size)]

    def process_blocks(self, blocks):
        """
        处理数据块并去重
        :param blocks: 输入块列表
        :return: (唯一块列表, 重删率)
        """
        unique_blocks = []
        total_size = sum(len(b) for b in blocks)
        dedup_size = 0

        with ThreadPoolExecutor() as executor:  # 并行处理
            futures = []
            for block in blocks:
                future = executor.submit(self._process_single_block, block)
                futures.append(future)
            
            for future in futures:
                is_unique, block_size = future.result()
                if not is_unique:
                    dedup_size += block_size

        dedup_rate = dedup_size / total_size * 100 if total_size > 0 else 0
        LOGGER.info(f"Deduplication completed | "
                   f"Total blocks: {len(blocks)} | "
                   f"Unique blocks: {len(unique_blocks)} | "
                   f"Dedup rate: {dedup_rate:.2f}%")
        
        return unique_blocks, dedup_rate

    def _process_single_block(self, block):
        """处理单个数据块（内部方法）"""
        block_hash = hashlib.sha256(block).hexdigest()
        cursor = self.conn.cursor()

        # 检查块是否存在
        cursor.execute('SELECT ref_count FROM blocks WHERE hash=?', (block_hash,))
        result = cursor.fetchone()

        if result:
            # 已存在，更新引用计数
            new_count = result[0] + 1
            cursor.execute('UPDATE blocks SET ref_count=? WHERE hash=?', 
                          (new_count, block_hash))
            self.conn.commit()
            return (False, len(block))
        else:
            # 新块，插入数据库
            cursor.execute('INSERT INTO blocks (hash, size) VALUES (?, ?)',
                          (block_hash, len(block)))
            self.conn.commit()
            return (True, len(block))

# ======================
# 加密模块
# ======================
class EncryptionModule:
    """
    数据加密模块
    功能：
    1. 使用AES-256-CBC算法加密
    2. 自动生成随机IV并存储在输出数据中
    3. 符合PKCS#7填充标准
    
    使用示例：
    >>> key = os.urandom(32)  # 256-bit key
    >>> encryptor = EncryptionModule(key)
    >>> encrypted = encryptor.encrypt(b'sensitive_data')
    >>> decrypted = encryptor.decrypt(encrypted)
    """
    
    def __init__(self, key):
        """
        初始化加密模块
        :param key: 加密密钥（32字节用于AES-256）
        """
        if len(key) not in [16, 24, 32]:
            raise ValueError("Key must be 16/24/32 bytes long")
        self.key = key
        LOGGER.info(f"Initialized encryption module | Algorithm: AES-256-CBC")

    def encrypt(self, data):
        """
        加密数据
        :param data: 原始字节数据
        :return: IV(16字节) + 加密数据
        """
        iv = get_random_bytes(16)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        padded_data = pad(data, AES.block_size)
        encrypted = cipher.encrypt(padded_data)
        return iv + encrypted

    def decrypt(self, data):
        """
        解密数据
        :param data: IV + 加密数据
        :return: 原始字节数据
        """
        iv = data[:16]
        ciphertext = data[16:]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        decrypted = unpad(cipher.decrypt(ciphertext), AES.block_size)
        return decrypted